import torch
from torch import nn
from transformers import BertModel
from transformers import BertTokenizer
import os
from torch.utils.data import Dataset, DataLoader


# 下载的预训练文件路径
BERT_PATH = r'F:\dataset\models\bert-base-chinese'
# 加载分词器
tokenizer = BertTokenizer.from_pretrained(BERT_PATH)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
save_path = './bert_checkpoint'


class BertClassifier(nn.Module):
    def __init__(self):
        super(BertClassifier, self).__init__()
        self.bert = BertModel.from_pretrained(BERT_PATH)
        self.dropout = nn.Dropout(0.5)
        self.linear = nn.Linear(768, 10)
        self.relu = nn.ReLU()

    def forward(self, input_id, mask):
        _, pooled_output = self.bert(input_ids=input_id, attention_mask=mask, return_dict=False)
        dropout_output = self.dropout(pooled_output)
        linear_output = self.linear(dropout_output)
        final_layer = self.relu(linear_output)
        return final_layer


# 加载模型
model = BertClassifier()
model.load_state_dict(torch.load(os.path.join(save_path, 'best.pt')))
model = model.to(device)
model.eval()
# 读取标签
real_labels = []
with open(r'F:\dataset\THUCNews\THUCNews\data\class.txt', 'r') as f:
    for row in f.readlines():
        real_labels.append(row.strip())

while True:
    text = input('新闻标题：')
    bert_input = tokenizer(text, padding='max_length',
                           max_length=35,
                           truncation=True,
                           return_tensors="pt")
    input_ids = bert_input['input_ids'].to(device)
    masks = bert_input['attention_mask'].unsqueeze(1).to(device)
    output = model(input_ids, masks)
    print(output)
    pred = output.argmax(dim=1)
    print(real_labels[pred])